// Copyright 2021 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.

package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"sort"
	"strings"

	"github.com/alessio/shellescape"
)

func getPackagesToQuery() string {
	// First list all test and binary targets.
	infos, err := os.ReadDir("pkg")
	if err != nil {
		panic(err)
	}
	var packagesToQuery []string
	for _, info := range infos {
		// We don't want to query into pkg/ui because it only contains a
		// single Go test target at its root which will be included below.
		// Querying into its subdirectories is unneeded and causes a pull from `npm`.
		if !info.IsDir() || info.Name() == "ui" {
			continue
		}
		packagesToQuery = append(packagesToQuery, fmt.Sprintf("//pkg/%s/...", info.Name()))
	}
	packagesToQuery = append(packagesToQuery, "//pkg/ui:*")
	return strings.Join(packagesToQuery, "+")
}

func generateTestSuites() {
	cmd := exec.Command(
		"bazel", "query",
		fmt.Sprintf(`kind("((go|sh)_(binary|library|test|transition_binary|transition_test))", %s)`, getPackagesToQuery()),
		"--output=label_kind",
	)
	buf, err := cmd.Output()
	if err != nil {
		log.Printf("Could not query Bazel tests: got error %v", err)
		var cmderr *exec.ExitError
		if errors.As(err, &cmderr) {
			log.Printf("Got error output: %s", string(cmderr.Stderr))
		} else {
			log.Printf("Run `%s` to reproduce the failure", shellescape.QuoteCommand(cmd.Args))
		}
		os.Exit(1)
	}
	var goLabels, testLabels []string
	for _, line := range strings.Split(string(buf[:]), "\n") {
		fields := strings.Fields(line)
		if len(fields) != 3 {
			continue
		}
		kind := fields[0]
		label := fields[2]
		if kind == "go_library" && !strings.Contains(label, "TxnStateTransition") {
			goLabels = append(goLabels, label)
		} else if kind == "go_test" {
			testLabels = append(testLabels, label)
			goLabels = append(goLabels, label)
		} else if kind == "go_transition_test" {
			goLabels = append(goLabels, label)
		} else if kind == "sh_test" {
			testLabels = append(testLabels, label)
		} else if (kind == "go_binary" || kind == "go_transition_binary") && !strings.HasSuffix(label, "_gomock_prog_bin") && !strings.Contains(label, "TxnStateTransitions") {
			goLabels = append(goLabels, label)
		}
	}
	sort.Strings(goLabels)
	sort.Strings(testLabels)

	f, err := os.Create("pkg/BUILD.bazel")
	if err != nil {
		log.Fatalf("Failed to open file `pkg/BUILD.bazel` - %v", err)
	}
	w := bufio.NewWriter(f)

	fmt.Fprintln(w, `# Code generated by generate-bazel-extra, DO NOT EDIT.
# gazelle:proto_strip_import_prefix /pkg

ALL_TESTS = [`)
	for _, label := range testLabels {
		fmt.Fprintf(w, "    %q,\n", label)
	}
	fmt.Fprintln(w, `]

GO_TARGETS = [`)
	for _, label := range goLabels {
		fmt.Fprintf(w, "    %q,\n", label)
	}
	fmt.Fprintln(w, "]")

	fmt.Fprintln(w, `
test_suite(
    name = "all_tests",
    tags = [
        "-integration",
    ],
    tests = ALL_TESTS,
)`)

	if err := w.Flush(); err != nil {
		log.Fatal(err)
	}
	if err := f.Close(); err != nil {
		log.Fatal(err)
	}
}

func main() {
	doTestSuites := flag.Bool("gen_test_suites", false, "generate test suites")
	flag.Parse()
	if *doTestSuites {
		generateTestSuites()
	}
}
